package com.ruoyi.common.utils;
import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.nio.file.Files;
import java.util.List;

/**
 * YOLO标注批量可视化工具
 * 功能：
 * - 批量处理指定文件夹内的图片和标注文件
 * - 自动创建输出目录
 * - 支持多图片格式（JPG/PNG/BMP）
 * - 不同类别显示不同颜色
 * - 自动处理坐标转换
 * - 错误处理和日志提示
 */
public class YOLOBatchVisualizer {

    // 类别颜色配置（按类别ID循环使用）
    private static final Color[] CLASS_COLORS = {
            new Color(255, 0, 0),   // 红色 - Class 0
            new Color(0, 255, 0),   // 绿色 - Class 1
            new Color(0, 0, 255),   // 蓝色 - Class 2
            new Color(255, 255, 0), // 黄色 - Class 3
            new Color(255, 0, 255)  // 品红 - Class 4
    };

    // 主要配置参数
    private static final String INPUT_DIR = "C:/Users/lenovo/Downloads/dataset/images/";   // 输入图片目录
    private static final String LABEL_DIR = "C:/Users/lenovo/Downloads/dataset/labels/";   // 标注文件目录
    private static final String OUTPUT_DIR = "C:/Users/lenovo/Downloads/dataset/output/";  // 输出目录
    private static final String OUTPUT_PREFIX = "vis_";          // 输出文件前缀

    public static void main(String[] args) {
        try {
            processAllImages();
            System.out.println("批量处理完成！");
        } catch (Exception e) {
            System.err.println("主程序错误: " + e.getMessage());
            e.printStackTrace();
        }
    }

    /**
     * 批量处理所有图片
     */
    private static void processAllImages() throws Exception {
        // 确保输出目录存在
        createDirectoryIfNotExists(OUTPUT_DIR);

        // 获取输入目录中的所有图片文件
        File[] imageFiles = new File(INPUT_DIR).listFiles((dir, name) ->
                name.matches(".*\\.(?i)(jpg|jpeg|png|bmp)$")
        );

        if (imageFiles == null || imageFiles.length == 0) {
            System.out.println("没有找到可处理的图片文件");
            return;
        }

        // 处理每个图片文件
        for (File imageFile : imageFiles) {
            try {
                String baseName = getFilenameWithoutExtension(imageFile.getName());
                processSingleImage(
                        imageFile.getAbsolutePath(),
                        LABEL_DIR + baseName + ".txt",
                        OUTPUT_DIR + OUTPUT_PREFIX + imageFile.getName()
                );
                System.out.println("已处理: " + imageFile.getName());
            } catch (Exception e) {
                System.err.printf("处理失败 [%s]: %s%n",
                        imageFile.getName(), e.getMessage());
            }
        }
    }

    /**
     * 处理单张图片
     * @param imagePath 图片文件路径
     * @param labelPath 标注文件路径
     * @param outputPath 输出文件路径
     */
    private static void processSingleImage(String imagePath,
                                           String labelPath,
                                           String outputPath) throws Exception {
        // 1. 加载图片
        BufferedImage image = ImageIO.read(new File(imagePath));
        if (image == null) {
            throw new IllegalArgumentException("无法读取图片文件");
        }

        // 2. 准备绘图上下文
        Graphics2D g2d = image.createGraphics();
        setupGraphicsQuality(g2d);

        // 3. 处理标注文件
        File labelFile = new File(labelPath);
        if (!labelFile.exists()) {
            g2d.dispose();
            throw new IllegalArgumentException("标注文件不存在");
        }

        List<String> lines = Files.readAllLines(labelFile.toPath());
        for (String line : lines) {
            drawAnnotation(line, image.getWidth(), image.getHeight(), g2d);
        }

        // 4. 保存结果
        g2d.dispose();
        String formatName = getFileExtension(outputPath);
        ImageIO.write(image, formatName, new File(outputPath));
    }

    /**
     * 绘制单个标注
     * @param annotation 标注行内容
     * @param imgWidth 图片宽度
     * @param imgHeight 图片高度
     * @param g2d 图形上下文
     */
    private static void drawAnnotation(String annotation,
                                       int imgWidth,
                                       int imgHeight,
                                       Graphics2D g2d) {
        try {
            String[] parts = annotation.trim().split("\\s+");
            if (parts.length < 5) return;

            // 解析数据
            int classId = Integer.parseInt(parts[0]);
            double xCenter = Double.parseDouble(parts[1]);
            double yCenter = Double.parseDouble(parts[2]);
            double width = Double.parseDouble(parts[3]);
            double height = Double.parseDouble(parts[4]);
            float confidence = (parts.length >= 6) ?
                    Float.parseFloat(parts[5]) : 1.0f;

            // 坐标转换
            Rectangle box = convertYoloToPixel(
                    xCenter, yCenter, width, height,
                    imgWidth, imgHeight
            );

            // 绘制图形
            Color color = CLASS_COLORS[classId % CLASS_COLORS.length];
            String label = String.format("C:%d (%.2f)", classId, confidence);
            drawBoundingBox(g2d, box, color, 10, label);
        } catch (NumberFormatException e) {
            System.err.println("无效标注行: " + annotation);
        }
    }

    /**
     * YOLO格式转像素坐标
     */
    private static Rectangle convertYoloToPixel(double xCenter, double yCenter,
                                                double width, double height,
                                                int imgW, int imgH) {
        int pixelX = (int) Math.round(xCenter * imgW - (width * imgW) / 2);
        int pixelY = (int) Math.round(yCenter * imgH - (height * imgH) / 2);
        int pixelWidth = (int) Math.round(width * imgW);
        int pixelHeight = (int) Math.round(height * imgH);

        return new Rectangle(
                Math.max(pixelX, 0),
                Math.max(pixelY, 0),
                Math.min(pixelWidth, imgW - pixelX),
                Math.min(pixelHeight, imgH - pixelY)
        );
    }

    /**
     * 绘制边界框和标签
     */
    private static void drawBoundingBox(Graphics2D g2d, Rectangle box,
                                        Color color, int thickness,
                                        String label) {
        // 保存原始样式
        Stroke originalStroke = g2d.getStroke();
        Color originalColor = g2d.getColor();

        // 设置新样式
        g2d.setColor(color);
        g2d.setStroke(new BasicStroke(thickness));

        // 绘制边界框
        g2d.drawRect(box.x, box.y, box.width, box.height);

        // 绘制标签背景
        FontMetrics fm = g2d.getFontMetrics();
        int textWidth = fm.stringWidth(label);
        int textHeight = fm.getHeight();
        g2d.fillRect(box.x, box.y - textHeight, textWidth + 4, textHeight);

        // 绘制标签文字
        g2d.setColor(Color.WHITE);
        g2d.drawString(label, box.x + 2, box.y - 2);

        // 恢复原始样式
        g2d.setStroke(originalStroke);
        g2d.setColor(originalColor);
    }

    /**
     * 初始化图形质量设置
     */
    private static void setupGraphicsQuality(Graphics2D g2d) {
        g2d.setRenderingHint(RenderingHints.KEY_ANTIALIASING,
                RenderingHints.VALUE_ANTIALIAS_ON);
        g2d.setRenderingHint(RenderingHints.KEY_TEXT_ANTIALIASING,
                RenderingHints.VALUE_TEXT_ANTIALIAS_ON);
        g2d.setFont(new Font("Arial", Font.BOLD, 14));
    }

    /**
     * 工具方法：创建目录（如果不存在）
     */
    private static void createDirectoryIfNotExists(String path) {
        File dir = new File(path);
        if (!dir.exists() && !dir.mkdirs()) {
            throw new RuntimeException("无法创建目录: " + path);
        }
    }

    /**
     * 工具方法：获取不带扩展名的文件名
     */
    private static String getFilenameWithoutExtension(String filename) {
        int dotIndex = filename.lastIndexOf('.');
        return (dotIndex == -1) ? filename : filename.substring(0, dotIndex);
    }

    /**
     * 工具方法：获取文件扩展名
     */
    private static String getFileExtension(String path) {
        int dotIndex = path.lastIndexOf('.');
        return (dotIndex == -1) ? "jpg" : path.substring(dotIndex + 1);
    }
}